import pandas as pd
import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Dropout, Embedding, GlobalAveragePooling1D
from tensorflow.keras.optimizers import Adam

# 1. 读取数据集
train_data = pd.read_csv('dataset/train.csv')
test_data = pd.read_csv('dataset/test.csv')

# 2. 数据预处理
X_train = train_data.drop(['id', 'target'], axis=1).values
y_train = train_data['target'].values

# 3. 切分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)

# 4. 构建Transformer模型
def build_transformer_model(input_shape):
    inputs = Input(shape=input_shape)
    x = Embedding(input_dim=input_shape[0], output_dim=64)(inputs)  # 假设input_shape[0]为特征数目
    x = GlobalAveragePooling1D()(x)
    x = Dense(64, activation='relu')(x)
    x = Dropout(0.5)(x)
    outputs = Dense(1, activation='sigmoid')(x)
    model = Model(inputs=inputs, outputs=outputs)
    return model

# 创建模型实例
model = build_transformer_model(X_train.shape[1:])

# 编译模型
model.compile(optimizer=Adam(learning_rate=0.001), loss='binary_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(X_train, y_train, validation_data=(X_val, y_val), epochs=10, batch_size=32)

# 5. 预测结果
X_test = test_data.drop(['id'], axis=1).values
y_pred = model.predict(X_test)

# 四舍五入为0或1
y_pred = np.round(y_pred).astype(int)

# 6. 保存预测结果为CSV文件
pd.DataFrame({'id': test_data['id'], 'target': y_pred.flatten()}).to_csv('result/transformer.csv', index=False)
